﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Corecalc;
using System.Diagnostics;

namespace Sisal
{
    public static partial class Procedures
    {
        public static void InitCalls(this SisalNode node) {
            List<SisalNode> ns = node.GetGraph();

            foreach (SisalNode n in ns) {
                if (n is BuiltInFunctionCallNode) {

                    BuiltInFunctionCallNode mfn = n as BuiltInFunctionCallNode;
                    Func<SisalNode[], double> timeEstimator = Function.TimeEstimatorFor(mfn.Name);
                    n.ExecutionTimeCost = timeEstimator(mfn.InputEdges.ToArray());

                } else if (n is ParallelNode) {
                    throw new System.NotImplementedException();
                } else if (n is CompoundNode) {

                    foreach (SisalNode nn in n.InputEdges) {
                        nn.InitCalls();
                    }

                }
            }
        }

        public static void AssignCosts(this SisalNode g) {
            List<SisalNode> nodes = g.GetGraph();

            foreach (SisalNode node in nodes) {
                if (node is ParallelNode) {
                    throw new System.NotImplementedException("...");
                } else if (node is CompoundNode) {
                    CompoundNode cn = node as CompoundNode;
                    foreach (SisalNode sg in cn.Graphs)
                        sg.AssignCosts();

                    cn.ExecutionTimeCost = cn.ComputeExecutionTimeCost();
                }
            }

            //Prevent recursive calls to non-formula nodes from accumulating costs
            if (g.IsFormulaNode) {
                foreach (SisalNode node in nodes) {
                    if (node != g)
                        g.ExecutionTimeCost += node.ExecutionTimeCost;
                }
            }
        }

        public static void UnknownSarkarChiefCostSetter(SisalGraph g) {
            //Topologically sorted list (not reversed, due to the use of a support graph instead a dependency graph)
            //g.FormulaNodes is topologically sorted

            foreach (SisalNode formulaNode in g.FormulaNodes) {
                formulaNode.InitCalls();
                formulaNode.AssignCosts();
            }
        }

        public static Partition PartitionFunctions(SisalGraph g) {
            //NB: The formula nodes in g must be reverse topologically sorted - or just topologically in our case (support graph vs. dependency graph)

            foreach (SisalNode formulaNode in g.FormulaNodes) {
                formulaNode.Overhead = CommunicationCosts.TSched +
                                       CommunicationCosts.bytesToCycles(formulaNode.T_in) +
                                       CommunicationCosts.bytesToCycles(formulaNode.T_out);

                formulaNode.CP = formulaNode.ExecutionTimeCost +
                                           CommunicationCosts.bytesToCycles(formulaNode.Overhead);

                //Assign costs and CP to root to to be read in PartitionGraph
                SisalGraph.Root.ExecutionTimeCost += formulaNode.ExecutionTimeCost;
                SisalGraph.Root.CP += formulaNode.CP;
            }

            List<SisalNode> pn = new List<SisalNode>();
            Dictionary<SisalNode, Task> tp = new Dictionary<SisalNode, Task>();
 
            foreach (SisalNode formulaNode in g.FormulaNodes) {
                List<SisalNode> ppn = InitPrimitiveNodes(formulaNode);

                foreach (SisalNode n in ppn) {
                    Task task = new Task();
                    task.AddNode(n, true);
                    tp.Add(n, task);
                    SetTaskCosts(task);
                }

                pn.AddRange(ppn);
            }

            Partition partition = new Partition(pn, tp);
            Tuple<Partition, int> bestIter = PartitionGraph(SisalGraph.Root, partition);
            Tuple<Partition, int> bestPartition = PartitionGraph(SisalGraph.Root, partition, true, bestIter.Item2);

            return bestPartition.Item1;
        }

        public static List<SisalNode> InitPrimitiveNodes(SisalNode g) {
            List<SisalNode> ns = g.GetGraph();
            List<SisalNode> pn = new List<SisalNode>();

            foreach (SisalNode n in ns) {
                if (n.ExecutionTimeCost < CommunicationCosts.TMin) {
                    pn = pn.Union(new List<SisalNode> { n }).ToList();

                    if (n is BuiltInFunctionCallNode) {
                        BuiltInFunctionCallNode bfn = n as BuiltInFunctionCallNode;
                        bfn.Par = false;
                    } else if (n is ParallelNode) {
                        ParallelNode ppn = n as ParallelNode;
                        ppn.Size = (int)Math.Ceiling(ppn.F);
                    }
                } else if (n is SimpleNode) {
                    pn = pn.Union(new List<SisalNode> { n }).ToList();
                } else if (n is BuiltInFunctionCallNode) {
                    BuiltInFunctionCallNode bfn = n as BuiltInFunctionCallNode;
                    pn = pn.Union(new List<SisalNode> { n }).ToList();
                    bfn.Par = true;
                } else if (n is ParallelNode) {
                    //ParallelNode ppn = n as ParallelNode;
                    //InitPrimitiveNodes(n);
                    //ppn.Size = 1;
                    throw new System.NotImplementedException();
                } else {
                    //Compound Node
                    CompoundNode node = n as CompoundNode;
                    foreach (SisalNode n2 in node.Graphs) {
		                List<SisalNode> pn2 = InitPrimitiveNodes(n2);
                        pn = pn.Union(pn2).ToList();
                    }
                    pn = pn.Union(new List<SisalNode> { node }).ToList();
                }
            }
            
            return pn;
        }

        //Sets the cost of a task after an ACTUAL merge
        public static void SetTaskCosts(Task t) {
            if (t.IsSingleton && t.First is ParallelNode) {
                throw new System.NotImplementedException();
            } else {
                t.O = CommunicationCosts.TSched + t.T_in() + t.T_out();
                t.CP = t.T() + t.O;
            }
        }

        //Sets task cost of tasks originated from testing merge candidates
        public static void SetTempTaskCosts(Task t) {
            if (t.IsSingleton && t.First is ParallelNode) {
                throw new System.NotImplementedException();
            } else {
                t.O = CommunicationCosts.TSched + t.T_in_Temp() + t.T_out_Temp();
                t.CP = t.T() + t.O;
            }
        }

        public static Tuple<Partition, int> PartitionGraph(SisalRootNode g, Partition p, bool stop = false, int bestIndex = 0) {
            int bestIter = 0; 
            int curIter = 0; 
            double bestCost = p.CostFunction(g.ExecutionTimeCost);
            double overheadSum = p.Tasks.Sum(t => t.Q * t.O);
            double oldOverhead = 0.0;
            bool done = false;
            Task T_new = null;
            bool broke = false;

            while (p.Tasks.Distinct().Count() > 1) {
                var tasks = p.Tasks.Distinct();

                if (stop && curIter == bestIndex) {
                    broke = true;
                    break;
                }

                curIter = curIter + 1; done = false;
                Task T_a = p.Tasks.OrderBy(t => t.Q * t.O).Last();

                if (T_a.IsSingleton) {
                    oldOverhead = T_a.Q * T_a.O;

                    SisalNode n_a = T_a.GetNode(0);
                    if (n_a is BuiltInFunctionCallNode) {
                        BuiltInFunctionCallNode bfn = n_a as BuiltInFunctionCallNode;
                        if (bfn.Par) {
                            bfn.Par = false;
                            done = true;
                            T_new = T_a;
                        }
                    } else if (n_a is ParallelNode) {
                        ParallelNode pn = n_a as ParallelNode;
                        if (2 * pn.Size <= pn.F) {
                            pn.Size = 2 * pn.Size;
                            done = true;
                            T_new = T_a;
                        }
                    }
                } else if(T_a.IncludesAllNodesInNode(T_a.Parent)) {
                    oldOverhead = T_a.Q * T_a.O;
                    SisalNode n_b = T_a.Parent.Parent();
                    p.PN = p.PN.Union(new List<SisalNode> { n_b }).ToList();    //Refer to section 2.d.ii in PartitionGraph. Precedence in the set operators?
                    p.PN = p.PN.Except(T_a.Nodes).ToList();

                    T_new = new Task();
                    T_new.AddNode(n_b, true);
                    p.TP[n_b] = T_new;
                    done = true;

                    if (n_b is ParallelNode)
                        (n_b as ParallelNode).Size = 1;
                }

                if (!done) {
                    Dictionary<SisalNode, Task> localtp;
                    Dictionary<Task, HashSet<Task>> path;

                    SetupMerge(T_a, p, out localtp, out path);
                    //Partition p_clone = p.Clone();
                    List<Task> sameParentAsT_a = p.TP.Values.Where(t => t.Parent == T_a.Parent && t != T_a).ToList();

                    Task T_b = null;
                    double doesntMatter;
                    double bestCP = double.MaxValue;

                    List<Task> resultTasks = new List<Task>();

                    foreach (var T_b_candidate in sameParentAsT_a) {
                        Task resultTask;
                        MergeTasks(T_a, T_b_candidate, localtp, false, path, out resultTask, out doesntMatter);
                        SetTempTaskCosts(resultTask);
                        SetCPUntilRoot(g, resultTask, p);

                        resultTasks.Add(resultTask);

                        if (resultTask.CP < bestCP) {
                            bestCP = resultTask.CP;
                            T_b = T_b_candidate;
                        }
                    }

                    //Merge T_a with the best candidate for T_b
                    MergeTasks(T_a, T_b, localtp, true, path, out T_new, out oldOverhead);
                    p.UpdateTaskMapping(T_a, T_b, T_new);
                }

                SetTaskCosts(T_new);
                SetCPUntilRoot(g, T_new, p);

                overheadSum = overheadSum - oldOverhead + T_new.Q * T_new.O;
                int processors = CommunicationCosts.ProcessorCount;
                p.Cost = Math.Max(g.CP * processors / g.ExecutionTimeCost, 1 + overheadSum / g.ExecutionTimeCost);

                if (p.Cost <= bestCost) {
                    bestIter = curIter;
                    bestCost = p.Cost;
                }
            }

            if (broke) {
                return new Tuple<Partition, int>(p, curIter);
            } else {
                return new Tuple<Partition, int>(p, bestIter);
            }
        }

        private static void SetCPUntilRoot(SisalNode g, Task T_new, Partition p) {
            SisalNode gPrime = T_new.Parent;
            while (true) {
                double oldGPrimeCP = gPrime.CP;

                SetGraphCP(gPrime, p);
                if (gPrime == g || (gPrime.CP == oldGPrimeCP))
                    break;
                SisalNode nPrime = gPrime.Parent();

                double oldNPrimeCP = nPrime.CP;
                SetENCP(nPrime);
                gPrime = nPrime.Parent();
                if (nPrime.CP == oldNPrimeCP)
                    break;
            }
        }

        public static void SetupMerge(Task T_a, Partition p, out Dictionary<SisalNode, Task> ltp, out Dictionary<Task, HashSet<Task>> pth) {
            Dictionary<SisalNode, Task> localtp = new Dictionary<SisalNode, Task>();
            Action<SisalNode, Task> addTaskToLocalTp = (sn, t) => {
                sn.TempTask = t;
                sn.Task = t;
                localtp.Add(sn, t);
            };

            SisalNode g = T_a.Parent;
            List<SisalNode> nodes = g.GetGraph();
            foreach (SisalNode n in nodes) {
                if (p.PN.Contains(n)) {
                    addTaskToLocalTp(n, p.TP[n]);
                } else {
                    Task t = new Task();
                    t.AddNode(n, true);

                    addTaskToLocalTp(n, t);
                }
            }

            Dictionary<Task, HashSet<Task>> path = new Dictionary<Task, HashSet<Task>>();
            foreach (Task t_i in localtp.Values.Distinct()) {
                HashSet<Task> reachable = new HashSet<Task>();

                foreach (SisalNode node in t_i.Nodes) {
                    node.DFS(n => reachable.Add(n.Task));
                }

                path.Add(t_i, reachable);
            }

            ltp = localtp;
            pth = path;
        }

        public static void MergeTasks(Task t_a, Task t_b, Dictionary<SisalNode, Task> localtp, bool final, Dictionary<Task, HashSet<Task>> path, out Task t, out double overhead) {
            Debug.Assert(t_a.Parent == t_b.Parent);

            Task t_new = new Task();
            double oldOverhead = 0.0;

            foreach (Task t_i in localtp.Values.Distinct()) {
                bool b1 = path[t_a].Contains(t_i), b2 = path[t_b].Contains(t_i), b3 = path[t_i].Contains(t_a), b4 = path[t_i].Contains(t_b);

                if((b1 || b2) && (b3 || b4)) {
                    foreach (SisalNode node in t_i.Nodes) {
                        t_new.AddNode(node, final);
                    }

                    oldOverhead = oldOverhead + t_i.Q * t_i.O;
                }
            }

            t = t_new;
            overhead = oldOverhead;
        }

        public static void SetGraphCP(SisalNode g, Partition p) {
            //Compute new localtp - different from the localtp in SetupMerge and MergeTasks
            Dictionary<SisalNode, Task> localtp = new Dictionary<SisalNode, Task>();
            Action<SisalNode, Task> addTaskToLocalTp = (sn, t) => {
                sn.TempTask = t;
                sn.Task = t;
                localtp.Add(sn, t);
            };

            List<SisalNode> nodes = g.GetGraph();
            foreach (SisalNode n in nodes) {
                if (p.PN.Contains(n)) {
                    addTaskToLocalTp(n, p.TP[n]);
                } else {
                    Task t = new Task();
                    t.AddNode(n, true);

                    addTaskToLocalTp(n, t);
                }
            }

            Dictionary<Task, HashSet<Task>> prec = new Dictionary<Task, HashSet<Task>>();
            foreach (SisalNode n_a in nodes) {
                foreach (SisalNode n_b in nodes) {
                    if (!prec.ContainsKey(n_a.Task))
                        prec.Add(n_a.Task, new HashSet<Task>());

                    if (localtp[n_a] != localtp[n_b] && n_b.InputEdges.Contains(n_a)) {
                        prec[n_a.Task].Add(n_b.Task);
                    }
                }
            }

            List<Task> list = Procedures.TopSort(prec.Keys, t => prec[t]);

            //Order of the task in list "list"
            //Is it OK to assume that any task's start must be equal to the sum of "start" and "cp" of the
            //start lying directly behind it in the list?
            for (int j = 0; j < list.Count(); j++) {
                int i = j - 1;
                list[j].Start = Math.Max(0, j == 0 ? 0 : list[i].Start + list[i].CP);
            }

            g.CP = Math.Max(0, list.Max(t => t.Start + t.CP));
        }

        public static void SetENCP(SisalNode n) {
            if (n is ParallelNode) {
                (n as ParallelNode).CP = (n as ParallelNode).Graph.CP;
            } else if (n is CompoundNode) {
                CompoundNode cn = n as CompoundNode;
                cn.CP = cn.Graphs.Sum(g => g.F * g.CP); //Does g.CriticalPath represent the sum of the CP values of the entire graph or only its own?
            } else {
                var a = "?";
            }
        }

        public static IEnumerable<SisalNode> TopSort(Dictionary<SisalNode, List<SisalNode>> outgoing) {
            return Procedures.TopSort(outgoing.Keys, n => outgoing[n]);
        }

        public static List<A> TopSort<A>(IEnumerable<A> items, Func<A, IEnumerable<A>> outgoing) {
            List<A> sortedList = new List<A>();
            List<A> marked = new List<A>();
            List<A> tempMarked = new List<A>();

            var notDistinct = items.SelectMany(t => (t as Task).Nodes).OrderBy(x => x.GetType().ToString());


            Action<A> visit = null;
            visit = n => {
                if (tempMarked.Contains(n)) throw new Exception("Uh oh..!!");
                if (!marked.Contains(n)) {
                    tempMarked.Add(n);

                    var o = outgoing(n);
                    foreach (A m in o)
                        visit(m);

                    marked.Add(n);
                    tempMarked.Remove(n);
                    sortedList.Add(n);      //Should've been AddFirst, we reverse the list instead
                }
            };

            Queue<A> unmarked = new Queue<A>(items);
            while (unmarked.Count > 0) {
                A n = unmarked.Dequeue();
                visit(n);
            }

            sortedList.Reverse();
            return sortedList;
        }
    }

    public class Partition
    {
        public List<SisalNode> PN;
        public Dictionary<SisalNode, Task> TP;

        public double Cost { get; set; }

        public Partition(List<SisalNode> pn, Dictionary<SisalNode, Task> tp) {
            PN = pn;
            TP = tp;
        }

        public double T_total() {
            return TP.Values.Distinct().Sum(t => t.Q * (t.T() + t.O));
        }

        public double T_crit() {
            return PN.Sum(n => n.ExecutionTimeCost) + Tasks.Distinct().Sum(t => t.O);
        }

        public void UpdateTaskMapping(Task t_a, Task t_b, Task t_new) {
            //List<SisalNode> keysToModify = new List<SisalNode>();
            //foreach (var key in TP.Keys) {
            //    if (TP[key] == t_a || TP[key] == t_b) {
            //        keysToModify.Add(key);
            //    }
            //}
            //foreach (var key in keysToModify) {
            //    TP[key] = t_new;
            //}

            foreach (var node in t_new.Nodes) {
                TP[node] = t_new;
            }
        }

        public double CostFunction(double T_seq) {
            if (Cost == 0.0) {
                double P = CommunicationCosts.ProcessorCount;

                return Math.Max(T_crit(), T_total() / P) / (T_seq / P);
            } else {
                return Cost;
            }
        }

        public Partition Clone() {
            List<SisalNode> newPn = new List<SisalNode>();
            newPn.AddRange(PN);

            Dictionary<SisalNode, Task> newTp = TP.ToDictionary(kv => kv.Key, kv => kv.Value);
            return new Partition(newPn, newTp);
        }

        public IEnumerable<Task> Tasks {
            get { return TP.Values; }
        }
    }

    public class Task
    {
        public double O { get; set; }
        public double CP { get; set; }
        public double Start { get; set; }

        public List<SisalNode> Nodes { get; set; }

        public Task() {
            Nodes = new List<SisalNode>();
        }

        public SisalNode First {
            get { return Nodes[0]; }
        }

        public bool IsSingleton {
            get { return Nodes.Count == 1; }
        }

        //Q is one in a spreadsheet context
        public double Q {
            get { return 1.0; }
        }

        public double T() {
            return Nodes.Sum(n => n.ExecutionTimeCost);
        }

        public SisalNode Parent {
            //All nodes in a task have the same parent
            get { return Nodes.Count() == 0 ? null : Nodes[0].Parent(); }
        }

        public bool IncludesAllNodesInNode(SisalNode n) {
            var graph = n.GetGraph();
            var exceptNodes = graph.Except(Nodes);

            return exceptNodes.Count() == 0;
        }

        public double T_in() {
            double sum = 0.0;
            foreach (SisalNode node in Nodes) {
                foreach (SisalNode ie in node.InputEdges) {
                    var sameTask = ie.Task == this;
                    if (!sameTask) {
                        sum += CommunicationCosts.Rc(ie.T_out);
                    }
                }
            }
            return sum;


            //return Nodes.Sum(n => {
            //    return n.InputEdges.Sum(ie => {

            //        var sameTask = ie.Task == this;
            //        if (sameTask) {
            //            return 0.0;
            //        } else {
            //            return CommunicationCosts.Rc(ie.T_out);
            //        }

            //        //return ie.Task == this ? 0.0 : CommunicationCosts.Rc(ie.T_out);
            //    });
            //});
        }

        public double T_out() {
            return Nodes.Sum(n => {
                return n.OutputEdges.Sum(oe => {
                    return oe.Task == this ? 0.0 : CommunicationCosts.Wc(n.T_out);
                });
            });
        }

        public double T_in_Temp() {
            double sum = 0.0;
            foreach (SisalNode node in Nodes) {
                foreach (SisalNode ie in node.InputEdges) {
                    var sameTask = ie.TempTask == this;
                    if (!sameTask) {
                        sum += CommunicationCosts.Rc(ie.T_out);
                    }
                }
            }
            return sum;
        }

        public double T_out_Temp() {
            return Nodes.Sum(n => {
                return n.OutputEdges.Sum(oe => {
                    return oe.TempTask == this ? 0.0 : CommunicationCosts.Wc(n.T_out);
                });
            });
        }

        public void AddNode(SisalNode n, bool actualMerge) {
            if (actualMerge) {
                n.Task = this;

                var parents = Nodes.Select(nn => nn.Parent()).Distinct();
                if (parents.Count() > 0 && !parents.Contains(n.Parent())) {
                    throw new Exception("");
                }
            }

            n.TempTask = this;
            Nodes.Add(n);
        }

        public SisalNode GetNode(int i) {
            return Nodes[i];
        }

        public A FoldNodes<A>(A init, Func<A, SisalNode, A> f) {
            A a = init;
            foreach (SisalNode node in Nodes) {
                a = f(a, node);
            }
            return a;
        }

        public override string ToString() {
            return string.Format("Task [HashCode: {0}, Nodes: {1}, Start: {2}, O: {3}, CP: {4}]", this.GetHashCode(), Nodes.Count(), Start, O, CP);
        }
    }
}
